# analyzes python code for specific features used in the inference

import dis, new
import __builtin__

def operiter(codeobj):
  """
  Returns all operations in the bytecode, i.e. without grouping them into
  expressions
  """
  pc = 0                                # program counter
  while pc < len(codeobj.co_code):
    instr = ord(codeobj.co_code[pc])   # current instruction
    # read the data (if any) and move the program counter to the next instr.
    if instr >= dis.HAVE_ARGUMENT:
      data = ord(codeobj.co_code[pc+1]) \
             + (ord(codeobj.co_code[pc+2]) << 8)
      pc += 3
    else:
      data = None
      pc += 1
    opname = dis.opname[instr]
    yield (opname, data)

def expriter(codeobj):
  """
  Returns all expressions in the code wherever possible and return the rest
  as flat operations
  """
  # now scan through all the instructions
  pc = 0                                # program counter
  stack = []                            # simulated stack
  while pc < len(codeobj.co_code):
    instr = ord(codeobj.co_code[pc])   # current instruction
    # read the data (if any) and move the program counter to the next instr.
    if instr >= dis.HAVE_ARGUMENT:
      data = ord(codeobj.co_code[pc+1]) \
             + (ord(codeobj.co_code[pc+2]) << 8)
      pc += 3
    else:
      data = None
      pc += 1
    opname = dis.opname[instr]
    #
    # OPERATORS
    #
    # nullary operations simply append to the stack
    if opname in ('LOAD_CONST', 'LOAD_NAME', 'BUILD_MAP', 'LOAD_GLOBAL',
                  'LOAD_FAST', 'LOAD_CLOSURE', 'LOAD_DEREF'):
      if data is not None:
        stack.append((opname, data))
      else:
        stack.append((opname,))
    # unary operations take one item from the stack and put one back
    elif opname in ('UNARY_POSITIVE', 'UNARY_NEGATIVE', 'UNARY_NOT',
                    'UNARY_CONVERT', 'UNARY_INVERT', 'GET_ITER',
                    'SLICE+0',
                    # the following take data
                    'LOAD_ATTR'):
      # add dummy items on the stack
      if len(stack) < 1:
        stack.append(None)
      if data is not None:
        stack[-1] = (opname, stack[-1], data)
      else:
        stack[-1] = (opname, stack[-1])
    # binary operations remove two items from the stack and put one back
    elif opname in ('BINARY_POWER', 'BINARY_MULTIPLY',
                    'BINARY_DIVIDE', 'BINARY_FLOOR_DIVIDE',
                    'BINARY_TRUE_DIVIDE', 'BINARY_MODULO',
                    'BINARY_ADD', 'BINARY_SUBTRACT',
                    'BINARY_SUBSCR', 'BINARY_LSHIFT',
                    'BINARY_RSHIFT', 'BINARY_AND', 'BINARY_XOR',
                    'BINARY_OR', 'INPLACE_POWER', 'INPLACE_MULTIPLY',
                    'INPLACE_DIVIDE', 'INPLACE_FLOOR_DIVIDE',
                    'INPLACE_TRUE_DIVIDE', 'INPLACE_MODULO',
                    'INPLACE_ADD', 'INPLACE_SUBTRACT',
                    'INPLACE_LSHIFT', 'INPLACE_RSHIFT',
                    'INPLACE_AND', 'INPLACE_XOR', 'INPLACE_OR',
                    'SLICE+1', 'SLICE+2',
                    # the following take data
                    'COMPARE_OP'):
      # add dummy items on the stack
      while len(stack) < 2:
        stack.insert(0, None)
      if data is not None:
        stack[-2:] = [(opname, stack[-2], stack[-1], data)]
      else:
        stack[-2:] = [(opname, stack[-2], stack[-1])]
    # tertiary operators take 3 items and put one back on the stack
    elif opname in ('SLICE+3',):
      # add dummy items on the stack
      while len(stack) < 3:
        stack.insert(0, None)
      stack[-3:] = [(opname, stack[-3], stack[-2], stack[-1])]
    #
    ## SPECIAL OPERATIONS
    #
    elif opname in ('BUILD_TUPLE', 'BUILD_LIST'):
      # add dummy items on the stack
      while len(stack) < data:
        stack.insert(0, None)      
      stack[-data:] = [tuple([opname] + stack[-data:])]
    elif opname in ('CALL_FUNCTION', 'CALL_FUNCTION_VAR', 'CALL_FUNCTION_KW',
                    'CALL_FUNCTION_VAR_KW'):
      keyargs = data >> 8               # high byte
      otherargs = data % 256            # low byte
      if opname in ('CALL_FUNCTION_VAR', 'CALL_FUNCTION_KW'):
        otherargs += 1
      elif opname == 'CALL_FUNCTION_VAR_KW':
        otherargs += 2
      total = 1 + 2 * keyargs + otherargs
      # add dummy items on the stack
      while len(stack) < total:
        stack.insert(0, None)     
      # operation name, function name, number of keyword args,
      # keyword arguments, positional arguments, [var arg list], [var arg dict]
      stack[-total:] = [(opname, stack[-total], keyargs) \
                        + tuple(stack[-total+1:])]
    elif opname == 'BUILD_SLICE':
      # add dummy items on the stack
      while len(stack) < data:
        stack.insert(0, None)      
      stack[-data:] = [(opname,) + tuple(stack[-data:])]
    #
    ## STATEMENTS
    #
    else:
      # first empty the stack
      for item in stack:
        yield item
      stack = []
      if data is not None:
        yield (opname, data)
      else:
        yield (opname,)
  for item in stack:
    yield item

def unused_arguments(func):
  """
  A unused argument is one which is not used in the function at all. If there
  are unused arguments we will return the number of contiguous unused arguments
  from the right. Example
  def z(d,i):
    return Categorical(theta(d))
  z(d,i) has i as a unused argument, hence unused_arguments(z) => 1
  Note that if we has declared this function as z(i,d) then we wouldn't
  find any unused arguments.
  """
  codeobj = func.func_code              # extract the code object
  # we don't handle variable arguments, and there is no point if there
  # are no parameters of this function
  if codeobj.co_flags & (4 | 8) or not codeobj.co_argcount:
    return 0
  numparams = codeobj.co_argcount       # the number of parameters
  #
  # check for unused parameters
  #
  # initially assume that all parameters are unused
  used_pars = [False for i in range(numparams)]
  # now, go through the code
  for opname, data in operiter(codeobj):
    if opname == 'LOAD_FAST' and data < numparams:
      used_pars[data] = True
  # count the number of contiguous unused parameters
  numunused = 0
  for used in used_pars[::-1]:
    if used:
      break
    else:
      numunused += 1
  return numunused

def extract_load_ref(func, opname, val):
  """
  Extract the object being referred in a load statement
  For e.g. LOAD_GLOBAL 0 or LOAD_DEREF 1
  """
  def wrapper(arg):
    def dummy(): return arg
    return dummy
  
  if opname == 'LOAD_GLOBAL':
    # try in globals and then in builtins
    try:
      return func.func_globals[func.func_code.co_names[val]]
    except KeyError:
      return getattr(__builtin__, func.func_code.co_names[val])
  
  elif opname == 'LOAD_DEREF':
    # for derefs we'll use the wrapper and the dummy function
    # to get the a codeobj which returns its first closure element
    clos = func.func_closure[val]
    f = wrapper(None)
    newf = new.function(f.func_code, f.func_globals, f.func_name,
                        f.func_defaults, (clos,))
    return newf()

  else:
    return None
  
  
def indicator_function(func):
  """
  A function has an indicator function if the value of the function depends
  only on the value of another function evaluated on its arguments:
  For example
  def coord(i):
    return MNormal(mean(z(i)), var(z(i)))
  Here coord(i) depends only on the value of z(i)

  We will return the indicator function or None. NOTE: if there is an import
  statement in this function then we can't accurately determine the indicator
  function (as the runtime mapping may be different), in which case we also
  return None
  """
  codeobj = func.func_code              # extract the code object
  # we don't handle variable arguments, and there is no point if there
  # are no parameters of this function
  if codeobj.co_flags & (4 | 8) or not codeobj.co_argcount:
    return 0
  numparams = codeobj.co_argcount       # the number of parameters
  # finds all occurrences of a parameter in the given expression tree
  def find_parameter(tree, parent=None):
    if tree[0] == 'LOAD_FAST' and tree[1] < numparams:
      yield (tree, parent)
    else:
      for node in tree:
        if type(node) is tuple:
          for x in find_parameter(node, tree):
            yield x
  #
  # now look for an indicator function
  #
  # scan each expression in the code
  indicator_fn = None
  for expr in expriter(codeobj):
    if expr[0].startswith("IMPORT"):
      # An import statement would change the namespace of the function and
      # we wouldn't really know which function is the indicator function
      return None
    # look for an occurrence of the parameters in the expression
    for tree, parent in find_parameter(expr):
      # the parent better be a function call to the existing indicator
      # fn (if any) and there shouldn't be parameters called by dict
      # or list
      if parent[0] != 'CALL_FUNCTION':
        return None
      fncalled = extract_load_ref(func, parent[1][0], parent[1][1])
      
      if (fncalled is None) or \
             (indicator_fn is not None and indicator_fn != fncalled) \
             or parent[2] != 0 or len(parent) != (3+numparams):
        return None
      # ensure that the parameters are called in exactly the correct order
      for i,node in enumerate(parent[3:]):
        if node != ('LOAD_FAST', i):
          return None
      indicator_fn = fncalled

  # finds all occurrences of the indicator fn in the given expression tree
  def find_indicator(tree, parent=None):
    if (tree[0] in ('LOAD_GLOBAL', 'LOAD_DEREF'))\
       and (extract_load_ref(func, tree[0], tree[1]) is indicator_fn):
      yield (tree, parent)
    else:
      for node in tree:
        if type(node) is tuple:
          for x in find_indicator(node, tree):
            yield x

  # now, check that the indicator function is always used in the same manner
  for expr in expriter(codeobj):
    for tree, parent in find_indicator(expr):
      # the indicator function must always occur as a function call
      if (parent is None) or (parent[0] != 'CALL_FUNCTION') \
         or (extract_load_ref(func, parent[1][0], parent[1][1]) \
             != indicator_fn) \
         or parent[2] != 0 or len(parent) != (3+numparams):
        return None
      
      # the arguments of the indicator function must always be the
      # current function's arguments
      for i,node in enumerate(parent[3:]):
        if node != ('LOAD_FAST', i):
          return None
      
  # return the indicator function
  return indicator_fn

def indicator_bound(func, indfn):
  """
  If indfn is an indicator for func returns a list of bound functions,
  i.e. functions whose arguments are bound to the indicator's value
  
  For example

  def coord(i):
    return MNormal(mean(z(i)), var(z(i)))

  Here z is the indicator for coord and the bound functions are mean and var
  """
  # sanity check
  if indicator_function(func) is not indfn: return None

  # set of functions whose argument depends on the indicator
  bound_functions = set()
  
  codeobj = func.func_code              # extract the code object
  # we don't handle variable arguments, and there is no point if there
  # are no parameters of this function
  if codeobj.co_flags & (4 | 8) or not codeobj.co_argcount:
    return None

  # finds all occurrences of the call to the indicator fn in the expr. tree
  def find_indicator(tree, parent=None):
    if tree[0] == 'CALL_FUNCTION' \
           and (tree[1][0] in ('LOAD_GLOBAL', 'LOAD_DEREF'))\
       and (extract_load_ref(func, tree[1][0], tree[1][1]) is indfn):
      yield (tree, parent)
    else:
      for node in tree:
        if type(node) is tuple:
          for x in find_indicator(node, tree):
            yield x
  #
  # now look for an indicator function
  #
  # scan each expression in the code
  for expr in expriter(codeobj):

    if expr[0].startswith("IMPORT"):
      # An import statement would change the namespace of the function and
      # we wouldn't really know which function is the indicator function
      return None

    # look for an occurrence of the parameters in the expression
    for tree, parent in find_indicator(expr):
      # the parent better be a function call to the existing indicator
      # fn (if any) and there shouldn't be parameters called by dict
      # or list
      if (parent is None) or parent[0] != 'CALL_FUNCTION' or parent[2]\
             or (len(parent) != 4):
        return None
      
      fncalled = extract_load_ref(func, parent[1][0], parent[1][1])

      bound_functions.add(fncalled)

  # finally, return the bound functions
  return bound_functions


## NOTE I'm no longer using stmtiter, but I've left this code around to
## remind me how to analyze python code. expriter above is much simpler
def stmtiter(codeobj):
  """
  Iterates through a code object returning the instructions for each
  statement it encounters. A statement can be an assignment of a variable
  or a procedure call (a function call with return value ignored)
  Note: we only handle functions without loops, list comprehensions, try
  catch blocks and nested classes or functions
  """
  # first locate all the addresses corresponding to distinct line numbers
  line_addr = {}                   # for each addr gives the line number if any
  tab = codeobj.co_lnotab
  line = codeobj.co_firstlineno
  addr = 0
  for i in range(0, len(tab), 2):
    addr = addr + ord(tab[i])
    line = line + ord(tab[i+1])
    line_addr[addr] = line
  # now scan through all the instructions
  pc = 0                                # program counter
  stack = []                            # simulated stack
  while pc < len(codeobj.co_code):
    instr = ord(codeobj.co_code[pc])   # current instruction
    # read the data (if any) and move the program counter to the next instr.
    if instr >= dis.HAVE_ARGUMENT:
      data = ord(codeobj.co_code[pc+1]) \
             + (ord(codeobj.co_code[pc+2]) << 8)
      pc += 3
    else:
      data = None
      pc += 1
    opname = dis.opname[instr]
    #
    ## stack maintenance
    #
    if opname == 'ROT_TWO':
      stack[-2:] = [stack[-1], stack[-2]]
    elif opname == 'ROT_THREE':
      stack[-3:] = [stack[-1], stack[-3], stack[-2]]
    elif opname == 'ROT_FOUR':
      stack[-4:] = [stack[-1], stack[-4], stack[-3],
                         stack[-2]]
    #
    ## create references
    #
    # NOTE: technically creating a reference should be a new statement, but
    #       we ignore that for now
    #
    elif opname == 'DUP_TOP':
      stack.append(stack[-1])
    elif opname == 'UNPACK_SEQUENCE':
      tos = stack.pop()
      for i in range(data-1,-1,-1):
        stack.append(('BINARY_SUBSCR', tos, i))
    elif opname == 'DUP_TOPX':
      lastx = stack[-data:]
      for item in lastx:
        stack.append(item)
    #
    # OPERATORS
    #
    # nullary operations simply append to the stack
    elif opname in ('LOAD_CONST', 'LOAD_NAME', 'BUILD_MAP', 'LOAD_GLOBAL',
                    'LOAD_FAST', 'LOAD_CLOSURE', 'LOAD_DEREF'):
      if data is not None:
        stack.append((opname, data))
      else:
        stack.append((opname,))
    # unary operations take one item from the stack and put one back
    elif opname in ('UNARY_POSITIVE', 'UNARY_NEGATIVE', 'UNARY_NOT',
                    'UNARY_CONVERT', 'UNARY_INVERT', 'GET_ITER',
                    'SLICE+0',
                    # the following take data
                    'LOAD_ATTR', 'IMPORT_FROM'):
      if data is not None:
        stack[-1] = (opname, stack[-1], data)
      else:
        stack[-1] = (opname, stack[-1])
    # binary operations remove two items from the stack and put one back
    elif opname in ('BINARY_POWER', 'BINARY_MULTIPLY',
                    'BINARY_DIVIDE', 'BINARY_FLOOR_DIVIDE',
                    'BINARY_TRUE_DIVIDE', 'BINARY_MODULO',
                    'BINARY_ADD', 'BINARY_SUBTRACT',
                    'BINARY_SUBSCR', 'BINARY_LSHIFT',
                    'BINARY_RSHIFT', 'BINARY_AND', 'BINARY_XOR',
                    'BINARY_OR', 'INPLACE_POWER', 'INPLACE_MULTIPLY',
                    'INPLACE_DIVIDE', 'INPLACE_FLOOR_DIVIDE',
                    'INPLACE_TRUE_DIVIDE', 'INPLACE_MODULO',
                    'INPLACE_ADD', 'INPLACE_SUBTRACT',
                    'INPLACE_LSHIFT', 'INPLACE_RSHIFT',
                    'INPLACE_AND', 'INPLACE_XOR', 'INPLACE_OR',
                    'SLICE+1', 'SLICE+2',
                    # the following take data
                    'COMPARE_OP', 'IMPORT_NAME'):
      if data is not None:
        stack[-2:] = [(opname, stack[-2], stack[-1], data)]
      else:
        stack[-2:] = [(opname, stack[-2], stack[-1])]
    # tertiary operators take 3 items and put one back on the stack
    elif opname in ('SLICE+3',):
      stack[-3:] = [(opname, stack[-3], stack[-2], stack[-1])]
    #
    ## SPECIAL OPERATIONS
    #
    elif opname in ('BUILD_TUPLE', 'BUILD_LIST'):
      stack[-data:] = [tuple([opname] + stack[-data:])]
    elif opname in ('CALL_FUNCTION', 'CALL_FUNCTION_VAR', 'CALL_FUNCTION_KW',
                    'CALL_FUNCTION_VAR_KW'):
      keyargs = data >> 8               # high byte
      otherargs = data % 256              # low byte
      if opname in ('CALL_FUNCTION_VAR', 'CALL_FUNCTION_KW'):
        otherargs += 1
      elif opname == 'CALL_FUNCTION_VAR_KW':
        otherargs += 2
      total = 1 + 2 * keyargs + otherargs
      # operation name, function name, number of keyword args,
      # keyword arguments, positional arguments, [var arg list], [var arg dict]
      stack[-total:] = [(opname, stack[-total], keyargs) \
                        + tuple(stack[-total+1:])]
    elif opname == 'BUILD_SLICE':
      stack[-data:] = [(opname,) + tuple(stack[-data:])]
    #
    ## STATEMENTS ##
    #
    # a nullary statement removes nothing from the stack
    elif opname in ('PRINT_NEWLINE', 'BREAK_LOOP',
                    # the following have data
                    'CONTINUE_LOOP', 'DELETE_NAME', 'DELETE_GLOBAL',
                    'SET_LINENO'):
      if data is not None:
        yield (opname, data)
      else:
        yield (opname,)
    # a unary statement removes one item from the stack
    elif opname in ('POP_TOP', 'DELETE_SLICE+0', 'PRINT_ITEM',
                    'PRINT_NEWLINE_TO', 'RETURN_VALUE', 'YIELD_VALUE',
                    'IMPORT_STAR', 'STORE_NAME', 'DELETE_ATTR',
                    'STORE_GLOBAL', 'STORE_FAST', 'STORE_DEREF'):
      if data is not None:
        yield (opname, stack[-1], data)
      else:
        yield (opname, stack[-1])
      stack = stack[:-1]
    # binary statements remove two items from the stack
    elif opname in ('STORE_SLICE+0', 'DELETE_SLICE+1', 'DELETE_SLICE+2',
                    'DELETE_SUBSCR', 'PRINT_ITEM_TO', 'LIST_APPEND',
                    'STORE_ATTR'):
      if data is not None:
        yield (opname, stack[-2], stack[-1], data)
      else:
        yield (opname, stack[-2], stack[-1])
      stack = stack[:-2]
    # tertiary statements remove three items from the stack
    elif opname in ('STORE_SLICE+1','STORE_SLICE+2', 'DELETE_SLICE+3',
                    'STORE_SUBSCR', 'EXEC_STMT', 'BUILD_CLASS'):
      yield (opname, stack[-3], stack[-2], stack[-1])
      stack = stack[:-3]
    # quaternary statements remove four items from the stack
    elif opname in ('STORE_SLICE+1','STORE_SLICE+2'):
      yield (opname, stack[-3], stack[-2], stack[-1])
      stack = stack[:-3]
    #
    ## SPECIAL STATEMENTS
    #
    elif opname == 'RAISE_VARARGS':
      yield (opname,) + tuple(stack[-data:])
      stack = stack[:-data]
    #
    ## JUMPS
    #
    elif opname in ('JUMP_FORWARD', 'JUMP_IF_TRUE', 'JUMP_IF_FALSE',
                    'JUMP_ABSOLUTE', 'FOR_ITER'):
      raise 'if then and loop not supported'
    #
    ## unknown
    #
    # NOTE: we are not handling the following
    #    'PRINT_EXPR' -- this is only for interactive mode in any case
    #    'LOAD_LOCALS', 'POP_BLOCK', 'END_FINALLY', 'WITH_CLEANUP'
    #    'SETUP_LOOP', 'SETUP_EXCEPT', 'SETUP_FINALLY'
    #    'STORE_MAP', 'MAKE_FUNCTION' 'MAKE_CLOSURE' 'EXTENDED_ARG'
    else:
      raise "Unexpected opcode: " + opname
  # after exiting the loop the stack should be empty
  assert not stack, ("Stack", stack)
    
